package org.base.common.util;

import lombok.Data;
import org.apache.commons.collections.CollectionUtils;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.TreeSet;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.stream.Collectors;

@Data
public class TreeUtils {

    /**
     * 构建树
     *
     * @param nodes          待构建树的节点
     * @param topParentValue 根节点的值
     * @param parentFunc     获取父节点的值的方法
     * @param idFunc         获取id的方法
     * @param setChildren    设置子元素的方法
     * @param <E>
     * @param <O>
     * @return
     */
    public static <E, O> List<E> newTree(List<E> nodes, O topParentValue, Function<E, O> parentFunc, Function<E, O> idFunc, BiConsumer<E, List<E>> setChildren) {
        if (CollectionUtils.isEmpty(nodes)) {
            return null;
        }
        List<E> tops = getChildren(nodes, topParentValue, parentFunc);
        if (CollectionUtils.isNotEmpty(tops)) {
            for (E node : tops) {
                O parentValue = idFunc.apply(node);
                List<E> children = newTree(nodes, parentValue, parentFunc, idFunc, setChildren);
                setChildren.accept(node, children);
            }
        }
        return tops;
    }


    private static <E, O> List<E> getChildren(List<E> nodes, O parentValue, Function<E, O> parentFunc) {
        if (CollectionUtils.isEmpty(nodes)) {
            return null;
        }
        return nodes.stream().filter(e -> parentFunc.apply(e).equals(parentValue)).collect(Collectors.toList());
    }

    /**
     * 获取祖先节点
     */
    public static <E, O> List<E> getAncestor(List<E> nodes, O parentValue, O topParentValue, Function<E, O> parentFunc, Function<E, O> idFunc) {
        if (CollectionUtils.isEmpty(nodes) || parentValue == null || parentValue.equals(topParentValue)) {
            return null;
        }
        List<E> result = new ArrayList<>();
        for (E node : nodes) {
            if (idFunc.apply(node).equals(parentValue)) {
                result.add(node);
                if (!parentFunc.apply(node).equals(topParentValue)) {
                    final List<E> ancestor = getAncestor(nodes, parentFunc.apply(node), topParentValue, parentFunc, idFunc);
                    if (CollectionUtils.isNotEmpty(ancestor)) {
                        result.addAll(ancestor);
                    }
                }
                break;
            }
        }
        return result;
    }

    /**
     * 获取子孙节点
     */
    public static <E, O> List<E> getGrandchildren(List<E> nodes, O parentValue, Function<E, O> parentFunc, Function<E, O> idFunc) {
        if (CollectionUtils.isEmpty(nodes) || parentValue == null) {
            return null;
        }
        List<E> result = new ArrayList<>();
        final List<E> children = getChildren(nodes, parentValue, parentFunc);
        if (CollectionUtils.isNotEmpty(children)) {
            result.addAll(children);
            for (E child : children) {
                final List<E> grandchildren = getGrandchildren(nodes, idFunc.apply(child), parentFunc, idFunc);
                if (CollectionUtils.isNotEmpty(grandchildren)) {
                    result.addAll(grandchildren);
                }
            }
        }
        return result;
    }

    /**
     * 节点补全，如果节点是中间节点，则会遍历获取父节点和子孙节点
     *
     * @param nodes
     * @param id
     * @param parentValue
     * @param topParentValue
     * @param parentFunc
     * @param idFunc
     * @param <E>
     * @param <O>
     * @return
     */
    public static <E, O> List<E> getTreeList(List<E> nodes, O id, O parentValue, O topParentValue, Function<E, O> parentFunc, Function<E, O> idFunc) {
        List<E> result = new ArrayList<>();
        final List<E> ancestor = getAncestor(nodes, parentValue, topParentValue, parentFunc, idFunc);
        final List<E> grandchildren = getGrandchildren(nodes, id, parentFunc, idFunc);
        if (CollectionUtils.isNotEmpty(ancestor)) {
            result.addAll(ancestor);
        }
        if (CollectionUtils.isNotEmpty(grandchildren)) {
            result.addAll(grandchildren);
        }
        return result;
    }

    /**
     * 节点去重
     * @param nodes
     * @param deduplicationFunc
     * @param <E>
     * @param <O>
     * @return
     */
    public static <E, O extends Comparable> List<E> deduplicationNodes(List<E> nodes, Function<E, O> deduplicationFunc) {
        if (CollectionUtils.isEmpty(nodes) || deduplicationFunc == null) {
            return null;
        }
        final TreeSet<E> set = new TreeSet<E>(Comparator.comparing(deduplicationFunc));
        set.addAll(nodes);
        return new ArrayList<>(set);
    }

    public static <E, O extends Comparable> List<E> getDeduplicationTreeList(List<E> nodes, List<E> parentNodes, O topParentValue, Function<E, O> parentFunc, Function<E, O> idFunc) {
        if (CollectionUtils.isEmpty(parentNodes)) {
            return null;
        }
        List<E> result = new ArrayList<>();
        for (E parentNode : parentNodes) {
            final List<E> treeList = getTreeList(nodes, idFunc.apply(parentNode), parentFunc.apply(parentNode), topParentValue, parentFunc, idFunc);
            if (CollectionUtils.isNotEmpty(treeList)) {
                result.addAll(treeList);
            }
        }
        result.addAll(parentNodes);
        result = deduplicationNodes(result, idFunc);
        return result;
    }

    public static <E, O extends Comparable> List<E> getTree(List<E> nodes, List<E> parentNodes, O topParentValue, Function<E, O> parentFunc, Function<E, O> idFunc, BiConsumer<E, List<E>> setChildren) {
        final List<E> treeList = getDeduplicationTreeList(nodes, parentNodes, topParentValue, parentFunc, idFunc);
        if (CollectionUtils.isEmpty(treeList)) {
            return null;
        } else {
            return newTree(nodes, topParentValue, parentFunc, idFunc, setChildren);
        }
    }
}
